from __future__ import division

import sys, os.path, argparse, logging, itertools, time

import numpy, scipy.constants, scipy.optimize, scipy.stats, numpy.random.mtrand

import common, wlc, kinetics_opt, stats
from retraction_unfolding_filter import RetractionPeak, RetractionCurve, CurveStretch

logger = logging.getLogger("kinetics_opt_main")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)


def approx_mode(forces):
    use_mode = False
    if use_mode:
        return stats.gumbel_l_loc_scale(numpy.asarray(forces, dtype=float))[0]
    else:
        return numpy.mean(forces)

def standard_method(kbt, p, curve_stretches, peak_indices_list):
    kv2indices = CurveStretch.kv2indices(curve_stretches)
    if len(kv2indices) < 2:
        return 0, 0
    kv2forces = {}
    for curve_stretch, peaks_ahead_list in zip(curve_stretches, peak_indices_list):
        if len(peaks_ahead_list):
            kv = curve_stretch.spring_constant_n_per_m * curve_stretch.retraction_velocity_m_per_s
            forces = kv2forces.setdefault(kv, [])
            forces.extend([
                    curve_stretch.peak_for_peaks_ahead(peaks_ahead).rupture_force_n
                    for peaks_ahead in peaks_ahead_list])
    pairs = [
        (numpy.log(kv), approx_mode(forces), len(forces), numpy.std(forces))
        for (kv, forces) in kv2forces.items()]
    x, y, c, force_std = tuple(map(list, zip(*pairs)))
    slope, intercept, r, p, stderr = scipy.stats.linregress(
        [float(xx) for xx in x], [float(yy) for yy in y])
    indices = numpy.argsort(x)
    logger.info("slope=%r intercept=%r r=%r", slope, intercept, r)
    alpha = 1 / slope
    k0 = alpha / numpy.exp(intercept * alpha)
    return alpha * kbt, k0

def show_statistics(curve_stretch_playlist, curve_stretches):
    v2indices = CurveStretch.v2indices(curve_stretches)
    v_indices_list = list(v2indices.items())
    forces_list = CurveStretch.forces_list(
        [[curve_stretches[index] for index in indices]
         for _, indices in v_indices_list])
    forces_pn_list = list(map(
            lambda f: numpy.asarray(
                f, dtype=kinetics_opt.bestfloat) * 1e12, forces_list))
    logger.info(curve_stretch_playlist)
    for (v, _), forces_pn in sorted(zip(v_indices_list, forces_pn_list)):
        if len(forces_pn):
            logger.info(
                "v=%.1f(nm/s) Min(pN)=%.1f Max(pN)=%.1f Average(pN)=%.1f Median(pN)=%.1f Std(pN)=%.1f Count=%d",
                v * 1e9, numpy.min(forces_pn), numpy.max(forces_pn),
                numpy.mean(forces_pn), numpy.median(forces_pn),
                numpy.std(forces_pn), len(forces_pn))
    logger.info(
        "Curves: %d   Non-empty curves: %d   Peaks: %d",
        len(curve_stretches),
        sum(bool(len(curve_stretch)) for curve_stretch in curve_stretches),
        sum(len(curve_stretch) for curve_stretch in curve_stretches))
        
def fit_standard_method(t, p, curve_stretches, peak_indices_list):
    dx, k0 = standard_method(
        scipy.constants.k * t, p, curve_stretches, peak_indices_list)
    logger.info("Standard method: dx=%r(nm) k0=%r(s^{-1})", dx * 1e9, k0)
    return dx, k0

def extract_basic_peaks(curve_stretches, peak_indices_list, newtons_per_unit):
    basic_peaks_list = [
        [kinetics_opt.BasicPeak(
                kinetics_opt.Protocol(
                    curve_stretch.retraction_velocity_m_per_s / curve_stretch.peak_for_peaks_ahead(i).contour_length_m,
                    curve_stretch.spring_constant_n_per_m * curve_stretch.peak_for_peaks_ahead(i).contour_length_m,
                    int(numpy.around(curve_stretch.peak_for_peaks_ahead(
                                i).start_force_n / newtons_per_unit)), i + 1),
                int(numpy.around(curve_stretch.peak_for_peaks_ahead(
                            i).rupture_force_n / newtons_per_unit)))
         for i in peak_indices]
        for curve_stretch, peak_indices in zip(curve_stretches, peak_indices_list)]
    return sum(basic_peaks_list, [])

def extract_diffuse_peaks(
    curve_stretches, peak_indices_list, newtons_per_unit, upper_bound_units,
    diffuse_distribution, diffuse_half_support_units):
    if diffuse_distribution != "none":
        if diffuse_distribution == "normal":
            m = kinetics_opt.DiffusePeak.normal_from_basic_peak
        elif diffuse_distribution == "truncated_normal":
            m = kinetics_opt.DiffusePeak.truncated_normal_from_basic_peak
        else:
            raise Exception("Unknown distribution: %s." % diffuse_distribution)
        diffuse_peaks_list = [
            [m(kinetics_opt.BasicPeak(
                        kinetics_opt.Protocol(
                            curve_stretch.retraction_velocity_m_per_s / curve_stretch.peak_for_peaks_ahead(i).contour_length_m,
                            curve_stretch.spring_constant_n_per_m * curve_stretch.peak_for_peaks_ahead(i).contour_length_m,
                            int(numpy.around(curve_stretch.peak_for_peaks_ahead(i).start_force_n / newtons_per_unit)), i + 1),
                        int(numpy.around(curve_stretch.peak_for_peaks_ahead(i).rupture_force_n / newtons_per_unit))),
               curve_stretch.noise_stddev_n / newtons_per_unit,
               diffuse_half_support_units, upper_bound_units)
             for i in peak_indices]
            for curve_stretch, peak_indices in zip(
                curve_stretches, peak_indices_list)]
        diffuse_peaks = sum(diffuse_peaks_list, [])
    else:
        diffuse_peaks = []
    return diffuse_peaks

def fit_mle(
    temperature_K, p, upper_bound_units, newtons_per_unit,
    dx_lower_bound, dx_upper_bound, fixed_dx_list,
    diffuse_distribution, diffuse_k0_lower_bound, diffuse_k0_upper_bound,
    diffuse_half_support_units, curve_stretches, peak_indices_list_list):
    basic_peaks = extract_basic_peaks(
        curve_stretches, peak_indices_list_list, newtons_per_unit)
    fixed_dx_k0_lls_list = []
    for fixed_dx in fixed_dx_list:
        world_time = time.time()
        biased_k0, fixed_dx_lls = kinetics_opt.estimate_biased_k0(
            temperature_K, p, fixed_dx, upper_bound_units,
            newtons_per_unit, basic_peaks)
        fixed_dx_k0_lls_list.append((biased_k0, fixed_dx_lls))
        logger.info(
            "MLE[dx=%r(nm)]: biased_k0=%r(s^-1), unbiased_k0=%r(s^-1), ll=%r, world_time=%.3f(s)",
            fixed_dx * 1e9, biased_k0,
            kinetics_opt.unbias_k0(biased_k0, len(basic_peaks)),
            fixed_dx_lls.sum(), time.time() - world_time)
    world_time = time.time()
    dx, k0, lls = kinetics_opt.estimate_dx_k0(
        temperature_K, p, upper_bound_units, newtons_per_unit, basic_peaks,
        dx_lower_bound, dx_upper_bound)
    logger.info(
        "MLE: dx=%r(nm), k0=%r(s^-1), ll=%r, world_time=%.3f(s)",
        dx * 1e9, k0, lls.sum(), time.time() - world_time)
    diffuse_peaks = extract_diffuse_peaks(
        curve_stretches, peak_indices_list_list, newtons_per_unit, upper_bound_units,
        diffuse_distribution, diffuse_half_support_units)
    if diffuse_peaks:
        diffuse_fixed_dx_k0_lls_list = []
        for fixed_dx in fixed_dx_list:
            world_time = time.time()
            diffuse_fixed_dx_k0, diffuse_fixed_dx_lls = kinetics_opt.estimate_k0_diffuse(
                temperature_K, p, fixed_dx, upper_bound_units,
                newtons_per_unit, diffuse_peaks,
                diffuse_k0_lower_bound, diffuse_k0_upper_bound)
            diffuse_fixed_dx_k0_lls_list.append(
                (diffuse_fixed_dx_k0, diffuse_fixed_dx_lls))
            logger.info(
                "Diffuse MLE[dx=%r(nm)]: k0=%r(s^-1), ll=%r, world_time=%.3f(s)",
                fixed_dx * 1e9, diffuse_fixed_dx_k0,
                diffuse_fixed_dx_lls.sum(), time.time() - world_time)
        world_time = time.time()
        diffuse_dx, diffuse_k0, diffuse_lls = kinetics_opt.estimate_dx_k0_diffuse(
            temperature_K, p, upper_bound_units, newtons_per_unit, diffuse_peaks,
            dx_lower_bound, dx_upper_bound,
            diffuse_k0_lower_bound, diffuse_k0_upper_bound)
        logger.info(
            "Diffuse MLE: dx=%r(nm), k0=%r(s^-1), ll=%r, world_time=%.3f(s)",
            diffuse_dx * 1e9, diffuse_k0, diffuse_lls.sum(),
            time.time() - world_time)
    else:
        diffuse_fixed_dx_k0, diffuse_fixed_dx_k0_lls_list, diffuse_dx, diffuse_k0, diffuse_lls = 0, [], 0, 0, []
    return fixed_dx_k0_lls_list, dx, k0, lls, diffuse_fixed_dx_k0, diffuse_fixed_dx_k0_lls_list, diffuse_dx, diffuse_k0, diffuse_lls

def fit_shared_dx_mle(
    temperature_K, p, upper_bound_units, newtons_per_unit,
    dx_lower_bound, dx_upper_bound,
    diffuse_distribution,
    diffuse_k0_lower_bound, diffuse_k0_upper_bound,
    diffuse_half_support_units, curve_stretches_list, peak_indices_list_list):
    common.check_equal(len(curve_stretches_list), len(peak_indices_list_list))
    basic_peaks_list = [
        extract_basic_peaks(curve_stretches, peak_indices_list, newtons_per_unit)
        for curve_stretches, peak_indices_list in zip(
            curve_stretches_list, peak_indices_list_list)]
    world_time = time.time()
    dx, k0_lls_list = kinetics_opt.estimate_shared_dx_k0(
        temperature_K, p, upper_bound_units, newtons_per_unit, basic_peaks_list,
        dx_lower_bound, dx_upper_bound)
    logger.info(
        "MLE: dx=%r(nm), k0=%r(s^-1), ll=%r, world_time=%.3f(s)",
        dx * 1e9, zip(*k0_lls_list)[0], [lls.sum() for _, lls in k0_lls_list],
        time.time() - world_time)
    if diffuse_distribution != "none":
        diffuse_peaks_list = [
            extract_diffuse_peaks(
                curve_stretches, peak_indices_list, newtons_per_unit, upper_bound_units,
                diffuse_distribution, diffuse_half_support_units)
            for curve_stretches, peak_indices_list in zip(
                curve_stretches_list, peak_indices_list_list)]
        world_time = time.time()
        diffuse_dx, diffuse_k0_lls_list = kinetics_opt.estimate_shared_dx_k0_diffuse(
            temperature_K, p, upper_bound_units, newtons_per_unit, diffuse_peaks_list,
            dx_lower_bound, dx_upper_bound,
            diffuse_k0_lower_bound, diffuse_k0_upper_bound)
        logger.info(
            "Diffuse MLE: dx=%r(nm), k0=%r(s^-1), ll=%r, world_time=%.3f(s)",
            diffuse_dx * 1e9, zip(*diffuse_k0_lls_list)[0],
            [diffuse_lls.sum() for _, diffuse_lls in diffuse_k0_lls_list],
            time.time() - world_time)
    else:
        diffuse_dx, diffuse_k0_lls_list = 0, []
    return dx, k0_lls_list, diffuse_dx, diffuse_k0_lls_list

def main(params):
    parser = argparse.ArgumentParser()
    parser.add_argument("--curve_stretch_playlists", nargs="+")
    parser.add_argument(
        "--fit_methods", nargs="*", choices=(
            "standard", "mle", "shared_dx_mle", "none"),
        default=["standard", "mle", "shared_dx_mle"])
    parser.add_argument("--temperature_K", type=float, default=301.15)
    parser.add_argument("--p_nm", type=float, default=wlc.default_p * 1e9)
    parser.add_argument("--upper_bound_units", type=int, default=1000)
    parser.add_argument("--pn_per_unit", type=float, default=1)
    parser.add_argument("--dx_lower_bound_nm", type=float, default=0.01)
    parser.add_argument("--dx_upper_bound_nm", type=float, default=1)
    parser.add_argument(
        "--fixed_dx_nm", type=float, nargs="+", default=[0.165, 0.25])
    parser.add_argument(
        "--diffuse_k0_lower_bound", type=float,
        default=kinetics_opt.default_diffuse_k0_lower_bound,
        help="only for diffuse")
    parser.add_argument(
        "--diffuse_k0_upper_bound", type=float,
        default=kinetics_opt.default_diffuse_k0_upper_bound,
        help="only for diffuse")
    # The diffuse parameters could in principle be different for each dataset,
    # but for now we just share them.
    parser.add_argument("--diffuse_half_support_units", type=int, default=0)
    # Since the stretches don't have the last peak, it's actually a good thing
    # even when we know the total number of modules, since the last peak may be an
    # unfinished one, i.e, still stretching thus underestimating the rupture force.
    parser.add_argument(
        "--fixed_modules", type=int, default=0,
        help="For synthetic data when the number of modules is fixed and known.")
    parser.add_argument(
        "--diffuse_distribution", choices=("normal", "truncated_normal", "none"),
        default="none")
    parser.add_argument(
        "--seeds", type=int, nargs="+", help="empty means no bootstrap")
    parser.add_argument("--results_repr")

    args = parser.parse_args(params)
    logger.info("%s", args)

    common.check_equal(len(args.curve_stretch_playlists), len(args.seeds))

    curve_stretches_list = []
    for curve_stretch_playlist in args.curve_stretch_playlists:
        with open(curve_stretch_playlist, "rb") as in_file:
            curve_stretches = [eval(line) for line in in_file]
        if args.fixed_modules:
            curve_stretches = [
                curve_stretch.with_unseen_peaks(
                    args.fixed_modules - len(curve_stretch))
                for curve_stretch in curve_stretches]
        show_statistics(curve_stretch_playlist, curve_stretches)
        curve_stretches_list.append(curve_stretches)

    peak_indices_list_list = []
    for curve_stretches, seed in zip(curve_stretches_list, args.seeds):
        unseen_peak_counts = [
            curve_stretch.unseen_peak_count() for curve_stretch in curve_stretches]
        peak_counts = [len(curve_stretch) for curve_stretch in curve_stretches]
        unchanged_peak_indices_list = [
            range(unseen_peak_count, peak_count)
            for unseen_peak_count, peak_count in zip(unseen_peak_counts, peak_counts)]
        if not seed:
            peak_indices_list = unchanged_peak_indices_list
        else:
            a = sum([[
                        (i, j) for j in range(unseen_peak_counts[i], peak_counts[i])]
                     for i in range(len(peak_counts))], [])
            r = numpy.random.mtrand.RandomState(seed)
            choices = r.choice(len(a), size=len(a), replace=True)
            peak_indices_list = [[] for i in range(len(peak_counts))]
            for k in choices:
                i, j = a[k]
                peak_indices_list[i].append(j)
            logger.debug("peak_indices_list=%s", peak_indices_list)
        peak_indices_list_list.append(peak_indices_list)

    p = args.p_nm / 1e9

    if "standard" in args.fit_methods:
        standard_method_dx_k0_list = [
            fit_standard_method(
                args.temperature_K, p, curve_stretches, peak_indices_list)
            for curve_stretches, peak_indices_list in zip(
                curve_stretches_list, peak_indices_list_list)]
    else:
        standard_method_dx_k0_list = []

    newtons_per_unit = args.pn_per_unit / 1e12
    dx_lower_bound = args.dx_lower_bound_nm / 1e9
    dx_upper_bound = args.dx_upper_bound_nm / 1e9
    mle_list = []
    if "mle" in args.fit_methods:
        fixed_dx = [fixed_dx_nm / 1e9 for fixed_dx_nm in args.fixed_dx_nm]
        for curve_stretches, peak_indices_list in zip(
            curve_stretches_list, peak_indices_list_list):
            mle_row = fit_mle(
                args.temperature_K, p, args.upper_bound_units, newtons_per_unit,
                dx_lower_bound, dx_upper_bound, fixed_dx,
                args.diffuse_distribution,
                args.diffuse_k0_lower_bound, args.diffuse_k0_upper_bound,
                args.diffuse_half_support_units, curve_stretches, peak_indices_list)
            mle_list.append(mle_row)
    if "shared_dx_mle" in args.fit_methods:
        shared_dx_mle_row = fit_shared_dx_mle(
                args.temperature_K, p, args.upper_bound_units, newtons_per_unit,
                dx_lower_bound, dx_upper_bound,
                args.diffuse_distribution,
                args.diffuse_k0_lower_bound, args.diffuse_k0_upper_bound,
                args.diffuse_half_support_units,
                curve_stretches_list, peak_indices_list_list)
    else:
        shared_dx_mle_row = ()
    if args.results_repr:
        po = numpy.get_printoptions()
        # http://en.wikipedia.org/wiki/Quadruple_precision
        numpy.set_printoptions(
            precision=36, threshold=numpy.inf, linewidth=numpy.inf)
        try :
            r = repr((
                    args.curve_stretch_playlists, args.seeds, peak_indices_list_list,
                    standard_method_dx_k0_list, mle_list, shared_dx_mle_row))
        finally:
            numpy.set_printoptions(
                precision=po["precision"],
                threshold=po["threshold"], linewidth=po["linewidth"])
        with open(args.results_repr, "wb") as out_file:
            out_file.write(r + "\n")
    return peak_indices_list_list, standard_method_dx_k0_list, mle_list, shared_dx_mle_row

if __name__ == "__main__":
    main(sys.argv[1:])
